from langchain_ollama import OllamaEmbeddings
from langchain_chroma import Chroma


def queryVectorDB(queryText):
    # 绑定向量模型
    Embeddings = OllamaEmbeddings(model="autumnzsd/nlp_gte_sentence-embedding_chinese-large:latest")

    db = Chroma(persist_directory="./vector_db", embedding_function=Embeddings)

    # 查询
    # retriever = db.as_retriever(search_kwargs={"k": 3})
    retriever = db.as_retriever(search_type="similarity_score_threshold",
                                search_kwargs={"k": 10, "score_threshold": 0.48})
    results = retriever.invoke(queryText)

    return "\n\n".join(doc.page_content for i, doc in enumerate(results, start=1))
